#!/user/bin/env python3 
# -*- coding: utf-8 -*-
"""
    __author__ = "wu" 
   Description :
"""


import tensorflow as tf
from tensorflow.keras import layers, models, datasets


def load_image(img_path, size=(32, 32)):

    if tf.strings.regex_full_match(img_path, ".*automobile.*"):
        label = tf.constant(1, tf.int8)
    else:
        label = tf.constant(0, tf.int8)

    img = tf.io.read_file(img_path)
    img = tf.image.decode_jpeg(img)
    img = tf.image.resize(img, size) / 255.0

    return (img, label)


def build_model_by_model():
     tf.keras.backend.clear_session()

     input = layers.Input(shape=(32, 32, 3))
     x = layers.Conv2D(32, kernel_size=(3, 3))(input)
     x = layers.MaxPool2D()(x)
     x = layers.Conv2D(64, kernel_size=(5, 5))(x)
     x = layers.MaxPool2D()(x)
     x = layers.Dropout(rate=0.1)(x)
     x = layers.Flatten()(x)
     x = layers.Dense(32, activation="relu")(x)
     output = layers.Dense(1, activation="sigmoid")(x)

     model = models.Model(inputs=input, outputs=output)

     model.summary()
     model.compile(optimizer=tf.keras.optimizers.Adam(learning_rate=0.001),
                   loss=tf.keras.losses.binary_crossentropy,
                   metrics=["accuracy"])

     return model


def main():

    ds_train = tf.data.Dataset.list_files("../data/cifar2/train/*/*.jpg")\
        .map(load_image, num_parallel_calls=tf.data.experimental.AUTOTUNE)\
        .shuffle(buffer_size=100).batch(8).prefetch(tf.data.experimental.AUTOTUNE)

    ds_test = tf.data.Dataset.list_files("../data/cifar2/test/*/*.jpg")\
        .map(load_image, num_parallel_calls=tf.data.experimental.AUTOTUNE)\
        .shuffle(buffer_size=100).batch(8).prefetch(tf.data.experimental.AUTOTUNE)

    model = build_model_by_model()
    history = model.fit(ds_train, epochs=10, validation_data=ds_test)


if __name__ == '__main__':
    main()

